# -*- coding: utf-8 -*-
import numpy as np
from sklearn import metrics as mr


def MIBIFTrain(data_x, data_y, select_feature_num):
    """
    Mutual Information-based Best Individual Feature, MIBIF
    得到互信息最大的n个特征及其配对特征的位置
    ----------
    data_x: L×B×2m   空间滤波后的calibration session数据
           m: CSP维数  B: 目标滤波段数  L: 训练数据trial总数
    data_y: 1D L  L 个trial对应的标签
        n: 选择互信息最大的前n个
    ----------
    feature_set: type'set'  互信息最大的n个特征及其成对特征的位置 n~2n 个
                坐标(CSP维度, 滤波带序号)
    """
    feature_len = data_x.shape[2]
    filter_num = data_x.shape[1]
    mutual_info = np.zeros([feature_len, filter_num])
    for i in range(filter_num):
        for j in range(feature_len):
            mutual_info[j, i] = mr.mutual_info_score(data_y, data_x[:, i, j])
    # mutual_info = MutualInfoScore(data_x, data_y)
    mutual_info_1d = mutual_info.reshape(-1)  # 降到一维
    pos = np.argsort(-mutual_info_1d)  # 降序排列 返回位置序列
    # 选择互信息最大的前n个，若与已选特征成对的特征未被选入则加入
    feature_set = set()
    post_list = []
    for i in range(select_feature_num):
        post = divmod(pos[i], filter_num)
        post_list.append(post)
    for i in post_list:
        feature_set.add(i)
        feature_set.add((feature_len-1-i[0], i[1]))
    return feature_set


def MIBIFSelect(data_x, feature_set):
    """
    Mutual Information-based Best Individual Feature, MIBIF
    根据位置筛选对应的特征，得到新的训练数据
    ----------
    data_x: L×B×2m (或单个trial 1D 2m×B)  空间滤波后的数据
           m: CSP维数  B: 目标滤波段数  L: 训练数据trial总数
    feature_set: type'set'  互信息最大的n个特征及其成对特征的位置 n~2n 个
    ----------
    feature:  (n~2n)×L (或单个trial 1D n~2n)
    """
    feature_len = len(feature_set)
    if len(data_x.shape) == 3:
        trial_size = data_x.shape[0]
        feature = np.zeros([trial_size, feature_len])
        for i in range(trial_size):
            index = 0
            for j in feature_set:
                feature[i, index] = data_x[i, j[1], j[0]]
                index = index + 1
    else:
        feature = np.zeros(feature_len)
        index = 0
        for j in feature_set:
            feature[index] = data_x[j[1], j[0]]
            index = index + 1
    return feature
